from sqlalchemy import Column, Integer, String, Index
from sqlalchemy.orm.session import Session

from airflow.models.base import Base
from airflow.utils.session import provide_session


class TaskFile(Base):
    """the class to get Task info."""

    __tablename__ = "idata_task_file"

    id = Column(Integer, primary_key=True)
    task_id = Column(Integer)
    file_path = Column(String(512))

    __table_args__ = (
        Index('file_task_id', task_id, file_path),
        Index('task_id', task_id),
        Index('file_path', file_path),
    )

    @staticmethod
    @provide_session
    def get_files(task_id, session: Session = None):
        records = session.query(TaskFile).filter(TaskFile.task_id == task_id).all()
        return [r.file_path for r in records]

    @staticmethod
    @provide_session
    def query(task_id=None, file_path=None, session: Session = None):
        qry = session.query(TaskFile)
        if task_id:
            qry = qry.filter(TaskFile.task_id == task_id)
        if file_path:
            qry = qry.filter(TaskFile.file_path == file_path)
        return qry.all()

    @provide_session
    def delete(self, session: Session = None):
        session.query(TaskFile).filter(TaskFile.task_id == self.id, TaskFile.file_path == self.file_path).delete()

    @provide_session
    def insert(self, session: Session = None):
        session.add(self)
